
package chen.yi.component.log.utils;

import cn.hutool.core.util.StrUtil;
import com.alibaba.ttl.TransmittableThreadLocal;
import com.alibaba.ttl.TtlCallable;
import com.alibaba.ttl.TtlRunnable;
import jakarta.servlet.http.HttpServletRequest;
import org.slf4j.MDC;

import java.util.Map;
import java.util.Objects;
import java.util.UUID;
import java.util.concurrent.Callable;

import static chen.yi.component.log.constant.TraceIdConstant.*;

/**
 * MDC工具类
 */
public final class ThreadMdcUtil {
    private static final TransmittableThreadLocal<Map<String, String>> context =
            new TransmittableThreadLocal<>();

    /**
     * 生成traceId
     * @param request 请求
     * @return traceId字符串
     */
    public static String generateTraceId(HttpServletRequest request) {
        String traceId = request != null ? request.getHeader(TRACE_ID_HEADER) : null;
        return StrUtil.isNotBlank(traceId) ? traceId :
                "TID-" + System.currentTimeMillis() + "-" +
                        UUID.randomUUID().toString().replace("-", "").substring(0, 8);
    }

    /**
     * 获取traceId
     * @return traceId字符串
     */
    public static String getTraceId(){
        return context.get().get(TRACE_ID);
    }
    // 安全清除上下文
    public static Map<String,String> get() {
        return context.get();
    }

    // 安全清除上下文
    public static void clear() {
        try {
            MDC.clear();
            context.remove();
        } catch (Exception e) {
            // 防止清除操作影响主流程
        }
    }

    // 带校验的上下文设置
    public static void setContextMap(Map<String, String> map) {
        if (map != null && !map.isEmpty()) {
            context.set(map);
            MDC.setContextMap(map);
        }
    }

    // 线程安全的包装方法
    public static Runnable wrap(Runnable runnable) {
        Objects.requireNonNull(runnable);
        Map<String, String> parentContext = MDC.getCopyOfContextMap();
        return TtlRunnable.get(() -> {
            try {
                MDC.setContextMap(parentContext);
                runnable.run();
            } finally {
                clear();
            }
        });
    }

    // 支持返回值的包装方法
    public static <T> Callable<T> wrap(Callable<T> callable) {
        Objects.requireNonNull(callable);
        Map<String, String> parentContext = MDC.getCopyOfContextMap();
        return TtlCallable.get(() -> {
            try {
                MDC.setContextMap(parentContext);
                return callable.call();
            } finally {
                clear();
            }
        });
    }


    // 上下文备份方法
    public static void backupAndSet() {
        Map<String, String> map = MDC.getCopyOfContextMap();
        if (map != null) {
            context.set(map);
        }
    }
}
